package cc.shacocloud.luanniao.core.service.ai.doubao

import com.volcengine.ark.runtime.service.ArkService
import org.springframework.ai.document.Document
import org.springframework.ai.document.MetadataMode
import org.springframework.ai.embedding.AbstractEmbeddingModel
import org.springframework.ai.embedding.EmbeddingRequest
import org.springframework.ai.embedding.EmbeddingResponse

/**
 * 火山方舟的豆包大模型
 *
 * https://www.volcengine.com/docs/82379/1330206
 * @author 思追(shaco)
 */
class DoubaoEmbeddingModel(
    val arkService: ArkService,
    val endpointId: String,
    val metadataMode: MetadataMode = MetadataMode.EMBED,
) : AbstractEmbeddingModel(), AutoCloseable {

    override fun call(request: EmbeddingRequest): EmbeddingResponse {
        val embeddingRequest = com.volcengine.ark.runtime.model.embeddings.EmbeddingRequest.builder()
            .model(endpointId)
            .input(request.instructions)
            .build()

        val result = arkService.createEmbeddings(embeddingRequest)
        return result.toEmbeddingResponse()
    }

    override fun embed(document: Document): FloatArray {
        val content = document.getFormattedContent(metadataMode)
        return embed(content)
    }

    override fun close() {
        arkService.shutdownExecutor()
    }

}